#include "Astar_searcher.h"
#include <cmath>
#include <cstdio>
using namespace std;
using namespace Eigen;

void AstarPathFinder::initGridMap(double _resolution, Vector3d global_xyz_l, Vector3d global_xyz_u, int max_x_id, int max_y_id, int max_z_id){
    gl_xl = global_xyz_l(0);
    gl_yl = global_xyz_l(1);
    gl_zl = global_xyz_l(2);

    gl_xu = global_xyz_u(0);
    gl_yu = global_xyz_u(1);
    gl_zu = global_xyz_u(2);
    
    GLX_SIZE = max_x_id;
    GLY_SIZE = max_y_id;
    GLZ_SIZE = max_z_id;
    GLYZ_SIZE  = GLY_SIZE * GLZ_SIZE;
    GLXYZ_SIZE = GLX_SIZE * GLYZ_SIZE;

    resolution = _resolution;
    inv_resolution = 1.0 / _resolution;    

    data = new uint8_t[GLXYZ_SIZE];
    memset(data, 0, GLXYZ_SIZE * sizeof(uint8_t));
    
    GridNodeMap = new GridNodePtr ** [GLX_SIZE];
    for(int i = 0; i < GLX_SIZE; i++){
        GridNodeMap[i] = new GridNodePtr * [GLY_SIZE];
        for(int j = 0; j < GLY_SIZE; j++){
            GridNodeMap[i][j] = new GridNodePtr [GLZ_SIZE];
            for( int k = 0; k < GLZ_SIZE;k++){
                Vector3i tmpIdx(i,j,k);
                Vector3d pos = gridIndex2coord(tmpIdx);
                GridNodeMap[i][j][k] = new GridNode(tmpIdx, pos);
            }
        }
    }
}

void AstarPathFinder::resetGrid(GridNodePtr ptr){
    ptr->id = 0;
    ptr->cameFrom = nullptr;
    ptr->gScore = inf;
    ptr->fScore = inf;
}

void AstarPathFinder::resetUsedGrids(){
    for(int i=0; i < GLX_SIZE ; i++)
        for(int j=0; j < GLY_SIZE ; j++)
            for(int k=0; k < GLZ_SIZE ; k++)
                resetGrid(GridNodeMap[i][j][k]);
}

void AstarPathFinder::setObs(const double coord_x, const double coord_y, const double coord_z){
    if( coord_x < gl_xl  || coord_y < gl_yl  || coord_z <  gl_zl || 
        coord_x >= gl_xu || coord_y >= gl_yu || coord_z >= gl_zu )
        return;

    int idx_x = static_cast<int>( (coord_x - gl_xl) * inv_resolution);
    int idx_y = static_cast<int>( (coord_y - gl_yl) * inv_resolution);
    int idx_z = static_cast<int>( (coord_z - gl_zl) * inv_resolution);      

    data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] = 1;
}

vector<Vector3d> AstarPathFinder::getVisitedNodes(){
    vector<Vector3d> visited_nodes;
    for(int i = 0; i < GLX_SIZE; i++)
        for(int j = 0; j < GLY_SIZE; j++)
            for(int k = 0; k < GLZ_SIZE; k++){   
                //if(GridNodeMap[i][j][k]->id != 0) // visualize all nodes in open and close list
                if(GridNodeMap[i][j][k]->id == -1)  // visualize nodes in close list only
                    visited_nodes.push_back(GridNodeMap[i][j][k]->coord);
            }

    ROS_WARN("visited_nodes size : %d", visited_nodes.size());
    return visited_nodes;
}

Vector3d AstarPathFinder::gridIndex2coord(const Vector3i & index){
    Vector3d pt;

    pt(0) = ((double)index(0) + 0.5) * resolution + gl_xl;
    pt(1) = ((double)index(1) + 0.5) * resolution + gl_yl;
    pt(2) = ((double)index(2) + 0.5) * resolution + gl_zl;

    return pt;
}

Vector3i AstarPathFinder::coord2gridIndex(const Vector3d & pt){
    Vector3i idx;
    idx <<  min( max( int( (pt(0) - gl_xl) * inv_resolution), 0), GLX_SIZE - 1),
            min( max( int( (pt(1) - gl_yl) * inv_resolution), 0), GLY_SIZE - 1),
            min( max( int( (pt(2) - gl_zl) * inv_resolution), 0), GLZ_SIZE - 1);                  
  
    return idx;
}

Eigen::Vector3d AstarPathFinder::coordRounding(const Eigen::Vector3d & coord)
{
    return gridIndex2coord(coord2gridIndex(coord));
}

inline bool AstarPathFinder::isOccupied(const Eigen::Vector3i & index) const
{
    return isOccupied(index(0), index(1), index(2));
}

inline bool AstarPathFinder::isFree(const Eigen::Vector3i & index) const
{
    return isFree(index(0), index(1), index(2));
}

inline bool AstarPathFinder::isOccupied(const int & idx_x, const int & idx_y, const int & idx_z) const 
{
    return  (idx_x >= 0 && idx_x < GLX_SIZE && idx_y >= 0 && idx_y < GLY_SIZE && idx_z >= 0 && idx_z < GLZ_SIZE && 
            (data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] == 1));
}

inline bool AstarPathFinder::isFree(const int & idx_x, const int & idx_y, const int & idx_z) const 
{
    return (idx_x >= 0 && idx_x < GLX_SIZE && idx_y >= 0 && idx_y < GLY_SIZE && idx_z >= 0 && idx_z < GLZ_SIZE && 
           (data[idx_x * GLYZ_SIZE + idx_y * GLZ_SIZE + idx_z] < 1));
}

inline void AstarPathFinder::AstarGetSucc(GridNodePtr currentPtr, vector<GridNodePtr> & neighborPtrSets, vector<double> & edgeCostSets)
{   
    neighborPtrSets.clear();
    edgeCostSets.clear();
    for(int dx = -1; dx <= 1; dx++){
        for (int dy = -1; dy <= 1; dy++){
            for (int dz = -1; dz <= 1; dz++){
                if (dx == 0 && dy == 0 && dz == 0)
                    continue;
                int x = currentPtr->index(0) + dx, y = currentPtr->index(1) + dy, z = currentPtr->index(2) + dz;
                if (x >= 0 && x < GLX_SIZE && y >= 0 && y < GLY_SIZE && z >= 0 && z < GLZ_SIZE){
                    if (GridNodeMap[x][y][z]->id != -1  && isFree(x, y, z)){
                        neighborPtrSets.push_back(GridNodeMap[x][y][z]);
                        edgeCostSets.push_back(sqrt(dx * dx + dy * dy + dz * dz));//记录实际距离代价g
                    }
                }
            }
        }

    }
}

double AstarPathFinder::getHeu(GridNodePtr node1, GridNodePtr node2){
    vector<int> dIndex;
    dIndex.push_back(abs(node1->index(0) - node2->index(0)));
    dIndex.push_back(abs(node1->index(1) - node2->index(1)));
    dIndex.push_back(abs(node1->index(2) - node2->index(2)));
    sort(dIndex.begin(), dIndex.end());
    double heuCost = M_SQRT3 * dIndex[0] + M_SQRT2 * (dIndex[1]-dIndex[0]) + dIndex[2] - dIndex[1] - dIndex[0];
    return heuCost * 1.00002;
}

void AstarPathFinder::AstarGraphSearch(Vector3d start_pt, Vector3d end_pt){
    ros::Time time_1 = ros::Time::now();    

    //index of start_point and end_point
    Vector3i start_idx = coord2gridIndex(start_pt);
    Vector3i end_idx   = coord2gridIndex(end_pt);
    goalIdx = end_idx;

    printf("Astar planner beginning ............");
    printf("start index is: %d, %d, %d\n", start_idx(0), start_idx(1), start_idx(2));
    printf("end index is: %d, %d, %d\n", end_idx(0), end_idx(1), end_idx(2));

    if (isOccupied(start_idx) || isOccupied(end_idx)) {
        ROS_WARN("start point or end point is occupied, astar find failed!!");
        return;
    }

    //position of start_point and end_point
    // move start_pt and end_pt to the center of the grid
    start_pt = gridIndex2coord(start_idx);
    end_pt   = gridIndex2coord(end_idx);

    //Initialize the pointers of struct GridNode which represent start node and goal node
    auto startPtr = GridNodeMap[start_idx(0)][start_idx(1)][start_idx(2)];
    auto endPtr   = GridNodeMap[end_idx(0)][end_idx(1)][end_idx(2)];


    //openSet is the open_list implemented through multimap in STL library
    openSet.clear();
    // currentPtr represents the node with lowest f(n) in the open_list
    GridNodePtr currentPtr  = nullptr;
    GridNodePtr neighborPtr = nullptr;

    //put start node in open set
    startPtr -> gScore = 0;
    startPtr -> fScore = getHeu(startPtr,endPtr);
    startPtr -> id = 1; 
    startPtr -> coord = start_pt;
    startPtr->nodeMapIt = openSet.insert( make_pair(startPtr -> fScore, startPtr) );
    /*
    STEP 2 :  some else preparatory works which should be done before while loop
    please write your code below
    */
    vector<GridNodePtr> neighborPtrSets;
    vector<double> edgeCostSets;

    // this is the main loop
    while ( !openSet.empty() ){
        // find the ptr which has the minimum f(n) value
        currentPtr = openSet.begin()->second;
        openSet.erase(openSet.begin());
        currentPtr->id = -1;

        // if the current node is the goal 
        if( currentPtr->index == goalIdx ){
            ros::Time time_2 = ros::Time::now();
            terminatePtr = currentPtr;
            ROS_WARN("[A*]{success}  Time in A*  is %f ms, path cost if %f m", (time_2 - time_1).toSec() * 1000.0, currentPtr->gScore * resolution );
            return;
        }
        //get the succession
        AstarGetSucc(currentPtr, neighborPtrSets, edgeCostSets);


        for(int i = 0; i < neighborPtrSets.size(); ++i){
            neighborPtr = neighborPtrSets[i];
            if(neighborPtr -> id == 0){
                neighborPtr -> gScore = currentPtr->gScore + edgeCostSets[i];
                neighborPtr -> fScore = neighborPtr-> gScore + getHeu(neighborPtr, endPtr);
                neighborPtr -> cameFrom = currentPtr;
                neighborPtr -> nodeMapIt = openSet.insert( make_pair(neighborPtr -> fScore, neighborPtr) );
                neighborPtr -> id = 1;
            }
            else if(neighborPtr->id == 1){
                if(currentPtr->gScore + edgeCostSets[i] < neighborPtr->gScore) {
                    neighborPtr->gScore = currentPtr->gScore + edgeCostSets[i];
                    neighborPtr->fScore = neighborPtr->gScore + getHeu(neighborPtr, endPtr);
                    neighborPtr->cameFrom = currentPtr;
                    openSet.erase(neighborPtr->nodeMapIt);
                    neighborPtr->nodeMapIt = openSet.insert(make_pair(neighborPtr->fScore, neighborPtr));
                }
            }
        }      
    }
    //if search fails
    ros::Time time_2 = ros::Time::now();
    if((time_2 - time_1).toSec() > 0.1)
        ROS_WARN("Time consume in Astar path finding is %f", (time_2 - time_1).toSec() );
}


vector<Vector3d> AstarPathFinder::getPath() 
{   
    vector<Vector3d> path;
    vector<GridNodePtr> gridPath;

    if(terminatePtr == nullptr){
        ROS_WARN("Get empty path!");
        return path;
    }else{
        auto tmp = terminatePtr;
        while(tmp->cameFrom != nullptr){
            gridPath.push_back(tmp);
            tmp = tmp->cameFrom;
        }
    }

    for (auto ptr: gridPath)
        path.push_back(ptr->coord);
        
    reverse(path.begin(),path.end());

    return path;
}